package ir

import vm.TacEmulator;
import ir.TacValueRef.Kind;


class TacAssembly {

	public static class AssembledProgram {
		public List<TacInstruction> code
		public Map<Object,Integer> labels
		public List<Integer> data
	}
	
	private static class Tokens {
		List<String[]> code
		List<String> data
	}
	
	def Tokens tokenize(String program) {
		def lines = program.readLines().iterator();
		Tokens tokens = new Tokens()
		tokens.code = lines.takeWhile { String line -> !(line =~ /^\s*\.data/) }
			.collect { String line ->
				def x = (line =~ /#.*$/).replaceAll("").trim()
				x ? x.split(/\s+/) : null
			}.findAll { it }
		tokens.data = lines.collect { it }
		tokens
	}
	
	def TacInstruction parseInstruction(String[] tuple) {
		def mo = null
		if (tuple.length == 0) {
			throw new RuntimeException("empty instruction") // should never happen!
		}
		if (tuple.length == 1 && (mo = tuple[0] =~ /^(:.*):?$/)) {
			new TacInstruction(":", parseValue(mo[0][1]));
		}
		else if (tuple.length > 1) {
			new TacInstruction(tuple[0], (TacValueRef[]) tuple[1..-1].collect{x -> parseValue x});
		}
		else
			throw new RuntimeException("invalid instruction: " + tuple)
	}
	
	def List<TacInstruction> parseDatum(String literal) {
		literal = literal.trim()
		def mo = null
		if ((mo = literal =~ /^(:[^\s#]+?)\s*(#.*)?$/)) {
			[new TacInstruction(":", parseValue(mo[0][1]))]
		}
		else if ((mo = literal =~ /^[(]\s*(:[^\s#]+)\s*[)]\s*(#.*)?$/)) {
			[new TacInstruction("#", parseValue(mo[0][1]))]
		}
		else if ((mo = literal =~ /^"(([^\\]|\\.)*)"\s*(#.*)?$/)) {
			mo[0][1].collect { new TacInstruction("#", new TacValueRef((int)it)) }
		}
		else if ((mo = literal =~ /^([+-]?\d+)\s*(#.*)?$/)) {
			[new TacInstruction("#", new TacValueRef(Integer.parseInt(mo[0][1])))]
		}
		else if (literal =~ /^(#.*)?$/) []
		else
			throw new RuntimeException("invalid data: " + literal)
	}
	
	def TacValueRef parseValue(String token) {
		def mo
		if (token =~ /^[+-]?\d+$/)
			new TacValueRef(Kind.IMMEDIATE, Integer.parseInt(token))
		else if ((mo = token =~ /^:(\d*?):?$/))  
			new TacValueRef(Kind.LABEL, Integer.parseInt(mo[0][1]))
		else if ((mo = token =~ /^:(.*?):?$/))  
			new TacValueRef(Kind.LABEL, mo[0][1])
		else if ((mo = token =~ /\$(\d+)/))  
			new TacValueRef(Kind.LOCAL, Integer.parseInt(mo[0][1]))
		else
			throw new RuntimeException("invalid token: " + token)
	}
	
	def AssembledProgram makeLabels(Iterable<TacInstruction> insn, int baseAddr=0) {
		def code = new ArrayList<TacInstruction>()
		def labels = new HashMap<Object,Integer>()
		insn.each { ins ->
			if (ins.op == ':') labels[{x -> x.name?x.name:x.val}(ins.args[0])] = code.size() + baseAddr;
			else code.add(ins)
		}
		new AssembledProgram(code:code, labels:labels)
	}
	
	def datum(AssembledProgram prog, TacValueRef val) {
		switch (val.kind) {
		case TacValueRef.Kind.LABEL:
			def loc = prog.labels.get(val.name?val.name:val.val); 
			if (loc != null) loc else throw new RuntimeException("no such label '" + val + "'")
			break
		case TacValueRef.Kind.IMMEDIATE: val.val; break
		default:
			throw new RuntimeException("invalid datum: " + val)
		}
	}
	
	def link(AssembledProgram code, AssembledProgram data) {
		code.labels.putAll data.labels
		code.data = data.code.collect { datum(code, it.getArguments()[0]) }
		code
	}
	
	def AssembledProgram assemble(Tokens program_tokens) {
		def insn =
			program_tokens.code.collect { tuple ->
				parseInstruction(tuple)
			}
		def data =
			program_tokens.data.collect { parseDatum it } flatten()
		link makeLabels(insn), makeLabels(data, 1)
	}
	
	def AssembledProgram assemble(String program_text) {
		assemble(tokenize(program_text))
	}
	
	static main(args) {
		/*
		 * This is just an example of how to use the assembler
		 */
		def prog = '''\
			:1
				>    $4 0 $2
				if=  $2 0 :2
				%    $1 $4 $5
				+    $5 0 $3 
				+    $4 0 $1
				+    $3 0 $4
				if=  0 0 :1
			:2
				if=  0 0 -1
			.data
			1
			"a"
			:str1
			2
			"y="
			(:2)
			(:str1)
		'''
		def input = [1:15l, 4:6l];
		
		def asm = new TacAssembly()
		def aprog = asm.assemble(prog)
		println aprog.data
		def (code, labels) = [aprog.code, aprog.labels]

		/* run emulator */
		def em = new TacEmulator()
		em.load(code, labels)
		em.load(input)
		em.go().dump();
		
	}

}
